!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief linres kernel functions
!> \par History
!>      created from qs_linres_methods
!> \author JGH
! **************************************************************************************************
MODULE qs_linres_kernel
   USE admm_types,                      ONLY: admm_type,&
                                              get_admm_env
   USE atomic_kind_types,               ONLY: atomic_kind_type,&
                                              get_atomic_kind
   USE cp_control_types,                ONLY: dft_control_type
   USE cp_dbcsr_api,                    ONLY: dbcsr_add,&
                                              dbcsr_copy,&
                                              dbcsr_create,&
                                              dbcsr_deallocate_matrix,&
                                              dbcsr_p_type,&
                                              dbcsr_set
   USE cp_dbcsr_operations,             ONLY: cp_dbcsr_sm_fm_multiply,&
                                              dbcsr_allocate_matrix_set,&
                                              dbcsr_deallocate_matrix_set
   USE cp_fm_types,                     ONLY: cp_fm_get_info,&
                                              cp_fm_type
   USE cp_log_handling,                 ONLY: cp_get_default_logger,&
                                              cp_logger_type,&
                                              cp_to_string
   USE hartree_local_methods,           ONLY: Vh_1c_gg_integrals
   USE hfx_energy_potential,            ONLY: integrate_four_center
   USE hfx_ri,                          ONLY: hfx_ri_update_ks
   USE hfx_types,                       ONLY: hfx_type
   USE input_constants,                 ONLY: do_admm_aux_exch_func_none,&
                                              do_admm_basis_projection,&
                                              do_admm_exch_scaling_none,&
                                              do_admm_purify_none,&
                                              kg_tnadd_embed
   USE input_section_types,             ONLY: section_get_ival,&
                                              section_get_lval,&
                                              section_get_rval,&
                                              section_vals_get,&
                                              section_vals_get_subs_vals,&
                                              section_vals_type,&
                                              section_vals_val_get
   USE kg_correction,                   ONLY: kg_ekin_subset
   USE kg_environment_types,            ONLY: kg_environment_type
   USE kinds,                           ONLY: default_string_length,&
                                              dp
   USE lri_environment_types,           ONLY: lri_density_type,&
                                              lri_environment_type,&
                                              lri_kind_type
   USE lri_ks_methods,                  ONLY: calculate_lri_ks_matrix
   USE message_passing,                 ONLY: mp_para_env_type
   USE mulliken,                        ONLY: ao_charges
   USE particle_types,                  ONLY: particle_type
   USE pw_env_types,                    ONLY: pw_env_get,&
                                              pw_env_type
   USE pw_methods,                      ONLY: pw_axpy,&
                                              pw_copy,&
                                              pw_scale,&
                                              pw_transfer
   USE pw_poisson_methods,              ONLY: pw_poisson_solve
   USE pw_poisson_types,                ONLY: pw_poisson_type
   USE pw_pool_types,                   ONLY: pw_pool_type
   USE pw_types,                        ONLY: pw_c1d_gs_type,&
                                              pw_r3d_rs_type
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE qs_fxc,                          ONLY: qs_fxc_analytic,&
                                              qs_fxc_fdiff
   USE qs_gapw_densities,               ONLY: prepare_gapw_den
   USE qs_integrate_potential,          ONLY: integrate_v_rspace,&
                                              integrate_v_rspace_diagonal,&
                                              integrate_v_rspace_one_center
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              get_qs_kind_set,&
                                              qs_kind_type
   USE qs_kpp1_env_types,               ONLY: qs_kpp1_env_type
   USE qs_ks_atom,                      ONLY: update_ks_atom
   USE qs_ks_types,                     ONLY: qs_ks_env_type
   USE qs_linres_types,                 ONLY: linres_control_type
   USE qs_neighbor_list_types,          ONLY: neighbor_list_set_p_type
   USE qs_p_env_methods,                ONLY: p_env_finish_kpp1
   USE qs_p_env_types,                  ONLY: qs_p_env_type
   USE qs_rho0_ggrid,                   ONLY: integrate_vhg0_rspace
   USE qs_rho_atom_types,               ONLY: rho_atom_type
   USE qs_rho_types,                    ONLY: qs_rho_get,&
                                              qs_rho_type
   USE qs_vxc_atom,                     ONLY: calculate_xc_2nd_deriv_atom
   USE task_list_types,                 ONLY: task_list_type
   USE xc,                              ONLY: xc_calc_2nd_deriv,&
                                              xc_prep_2nd_deriv
   USE xc_derivatives,                  ONLY: xc_functionals_get_needs
   USE xc_rho_cflags_types,             ONLY: xc_rho_cflags_type
   USE xc_rho_set_types,                ONLY: xc_rho_set_create,&
                                              xc_rho_set_release,&
                                              xc_rho_set_type,&
                                              xc_rho_set_update
   USE xtb_ehess,                       ONLY: xtb_coulomb_hessian
   USE xtb_types,                       ONLY: get_xtb_atom_param,&
                                              xtb_atom_type
#include "./base/base_uses.f90"

   IMPLICIT NONE

   PRIVATE

   ! *** Public subroutines ***
   PUBLIC :: apply_xc_admm
   PUBLIC :: apply_hfx
   PUBLIC :: apply_op_2
   PUBLIC :: hfx_matrix

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'qs_linres_kernel'

! **************************************************************************************************

CONTAINS

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param p_env ...
!> \param c0 ...
!> \param Av ...
! **************************************************************************************************
   SUBROUTINE apply_op_2(qs_env, p_env, c0, Av)
      !
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_p_env_type)                                :: p_env
      TYPE(cp_fm_type), DIMENSION(:), INTENT(IN)         :: c0
      TYPE(cp_fm_type), DIMENSION(:), INTENT(INOUT)      :: Av

      INTEGER                                            :: ispin, ncol
      TYPE(dft_control_type), POINTER                    :: dft_control

      CALL get_qs_env(qs_env=qs_env, dft_control=dft_control)
      IF (dft_control%qs_control%semi_empirical) THEN
         CPABORT("Linear response not available with SE methods")
      ELSEIF (dft_control%qs_control%dftb) THEN
         CPABORT("Linear response not available with DFTB")
      ELSEIF (dft_control%qs_control%xtb) THEN
         CALL apply_op_2_xtb(qs_env, p_env)
      ELSE
         CALL apply_op_2_dft(qs_env, p_env)
         CALL apply_hfx(qs_env, p_env)
         CALL apply_xc_admm(qs_env, p_env)
         IF (dft_control%do_admm) CALL p_env_finish_kpp1(qs_env, p_env)
      END IF

      DO ispin = 1, SIZE(c0)
         CALL cp_fm_get_info(c0(ispin), ncol_global=ncol)
         CALL cp_dbcsr_sm_fm_multiply(p_env%kpp1(ispin)%matrix, &
                                      c0(ispin), &
                                      Av(ispin), &
                                      ncol=ncol, alpha=1.0_dp, beta=1.0_dp)
      END DO

   END SUBROUTINE apply_op_2

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param p_env ...
! **************************************************************************************************
   SUBROUTINE apply_op_2_dft(qs_env, p_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_p_env_type)                                :: p_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'apply_op_2_dft'

      INTEGER                                            :: handle, ikind, ispin, nkind, ns, nspins
      LOGICAL                                            :: deriv2_analytic, gapw, gapw_xc, &
                                                            lr_triplet, lrigpw
      REAL(KIND=dp)                                      :: alpha, ekin_mol, energy_hartree, &
                                                            energy_hartree_1c
      TYPE(admm_type), POINTER                           :: admm_env
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(cp_logger_type), POINTER                      :: logger
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: k1mat, matrix_s, rho1_ao, rho_ao
      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: ksmat, psmat
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(kg_environment_type), POINTER                 :: kg_env
      TYPE(linres_control_type), POINTER                 :: linres_control
      TYPE(lri_density_type), POINTER                    :: lri_density
      TYPE(lri_environment_type), POINTER                :: lri_env
      TYPE(lri_kind_type), DIMENSION(:), POINTER         :: lri_v_int
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(pw_c1d_gs_type)                               :: rho1_tot_gspace, v_hartree_gspace
      TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER        :: rho1_g
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_poisson_type), POINTER                     :: poisson_env
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool
      TYPE(pw_r3d_rs_type)                               :: v_hartree_rspace
      TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER        :: rho1_r, rho_r, tau1_r, v_rspace_new, &
                                                            v_xc, v_xc_tau
      TYPE(qs_kpp1_env_type), POINTER                    :: kpp1_env
      TYPE(qs_ks_env_type), POINTER                      :: ks_env
      TYPE(qs_rho_type), POINTER                         :: rho, rho0, rho1, rho1_xc, rho1a, &
                                                            rho_aux, rho_xc
      TYPE(rho_atom_type), DIMENSION(:), POINTER         :: rho1_atom_set, rho_atom_set
      TYPE(section_vals_type), POINTER                   :: input, xc_section, xc_section_aux

      CALL timeset(routineN, handle)

      NULLIFY (auxbas_pw_pool, pw_env, v_rspace_new, para_env, rho1_r, &
               v_xc, rho1_ao, rho_ao, poisson_env, input, rho, dft_control, &
               logger, rho1_g, v_xc_tau)
      logger => cp_get_default_logger()

      energy_hartree = 0.0_dp
      energy_hartree_1c = 0.0_dp

      CPASSERT(ASSOCIATED(p_env%kpp1))
      CPASSERT(ASSOCIATED(p_env%kpp1_env))
      kpp1_env => p_env%kpp1_env

      CALL get_qs_env(qs_env=qs_env, &
                      ks_env=ks_env, &
                      pw_env=pw_env, &
                      input=input, &
                      admm_env=admm_env, &
                      para_env=para_env, &
                      rho=rho, &
                      rho_xc=rho_xc, &
                      linres_control=linres_control, &
                      dft_control=dft_control)

      gapw = dft_control%qs_control%gapw
      gapw_xc = dft_control%qs_control%gapw_xc
      lr_triplet = linres_control%lr_triplet

      rho1 => p_env%rho1
      rho1_xc => p_env%rho1_xc
      CPASSERT(ASSOCIATED(rho1))
      IF (gapw_xc) THEN
         CPASSERT(ASSOCIATED(rho1_xc))
      END IF

      CALL qs_rho_get(rho, rho_ao=rho_ao, rho_r=rho_r)
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)

      nspins = SIZE(p_env%kpp1)
      lrigpw = dft_control%qs_control%lrigpw
      IF (lrigpw) THEN
         CALL get_qs_env(qs_env, &
                         lri_env=lri_env, &
                         lri_density=lri_density, &
                         atomic_kind_set=atomic_kind_set)
      END IF

      IF (.NOT. ASSOCIATED(kpp1_env%v_ao)) THEN
         CALL get_qs_env(qs_env, matrix_s=matrix_s)
         CALL dbcsr_allocate_matrix_set(kpp1_env%v_ao, nspins)
         DO ispin = 1, nspins
            ALLOCATE (kpp1_env%v_ao(ispin)%matrix)
            CALL dbcsr_copy(kpp1_env%v_ao(ispin)%matrix, matrix_s(1)%matrix, &
                            name="kpp1%v_ao-"//ADJUSTL(cp_to_string(ispin)))
         END DO
      END IF

      IF (dft_control%do_admm) THEN
         xc_section => admm_env%xc_section_primary
      ELSE
         xc_section => section_vals_get_subs_vals(input, "DFT%XC")
      END IF

      ! gets the tmp grids
      CPASSERT(ASSOCIATED(pw_env))
      CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool, &
                      poisson_env=poisson_env)
      CALL auxbas_pw_pool%create_pw(v_hartree_gspace)
      CALL auxbas_pw_pool%create_pw(v_hartree_rspace)

      IF (gapw .OR. gapw_xc) &
         CALL prepare_gapw_den(qs_env, p_env%local_rho_set, do_rho0=(.NOT. gapw_xc))

      ! *** calculate the hartree potential on the total density ***
      CALL auxbas_pw_pool%create_pw(rho1_tot_gspace)

      CALL qs_rho_get(rho1, rho_g=rho1_g)
      CALL pw_copy(rho1_g(1), rho1_tot_gspace)
      DO ispin = 2, nspins
         CALL pw_axpy(rho1_g(ispin), rho1_tot_gspace)
      END DO
      IF (gapw) THEN
         CALL pw_axpy(p_env%local_rho_set%rho0_mpole%rho0_s_gs, rho1_tot_gspace)
         IF (ASSOCIATED(p_env%local_rho_set%rho0_mpole%rhoz_cneo_s_gs)) THEN
            CALL pw_axpy(p_env%local_rho_set%rho0_mpole%rhoz_cneo_s_gs, rho1_tot_gspace)
         END IF
      END IF

      IF (.NOT. (nspins == 1 .AND. lr_triplet)) THEN
         CALL pw_poisson_solve(poisson_env, rho1_tot_gspace, &
                               energy_hartree, &
                               v_hartree_gspace)
         CALL pw_transfer(v_hartree_gspace, v_hartree_rspace)
      END IF

      CALL auxbas_pw_pool%give_back_pw(rho1_tot_gspace)

      ! *** calculate the xc potential ***
      NULLIFY (rho1a)
      IF (gapw_xc) THEN
         rho0 => rho_xc
         rho1a => rho1_xc
      ELSE
         rho0 => rho
         rho1a => rho1
      END IF

      deriv2_analytic = section_get_lval(xc_section, "2ND_DERIV_ANALYTICAL")
      NULLIFY (v_xc_tau)
      IF (deriv2_analytic) THEN
         CALL qs_rho_get(rho1a, rho_r=rho1_r, tau_r=tau1_r)
         CALL qs_fxc_analytic(rho0, rho1_r, tau1_r, xc_section, auxbas_pw_pool, lr_triplet, v_xc, v_xc_tau)
         IF (gapw .OR. gapw_xc) THEN
            CALL get_qs_env(qs_env, rho_atom_set=rho_atom_set)
            rho1_atom_set => p_env%local_rho_set%rho_atom_set
            CALL calculate_xc_2nd_deriv_atom(rho_atom_set, rho1_atom_set, qs_env, xc_section, para_env, &
                                             do_triplet=lr_triplet)
         END IF
      ELSE
         CALL qs_fxc_fdiff(ks_env, rho0, rho1a, xc_section, 6, lr_triplet, v_xc, v_xc_tau)
         CPASSERT((.NOT. gapw) .AND. (.NOT. gapw_xc))
      END IF

      v_rspace_new => v_xc
      NULLIFY (v_xc)

      CALL pw_scale(v_hartree_rspace, v_hartree_rspace%pw_grid%dvol)
      DO ispin = 1, nspins
         CALL pw_scale(v_rspace_new(ispin), v_rspace_new(ispin)%pw_grid%dvol)
         IF (ASSOCIATED(v_xc_tau)) CALL pw_scale(v_xc_tau(ispin), v_xc_tau(ispin)%pw_grid%dvol)
      END DO

      ! ADMM Correction
      IF (dft_control%do_admm) THEN
         IF (admm_env%aux_exch_func /= do_admm_aux_exch_func_none) THEN
            IF (.NOT. ASSOCIATED(kpp1_env%deriv_set_admm)) THEN
               CPASSERT(.NOT. lr_triplet)
               xc_section_aux => admm_env%xc_section_aux
               CALL get_admm_env(qs_env%admm_env, rho_aux_fit=rho_aux)
               CALL qs_rho_get(rho_aux, rho_r=rho_r)
               ALLOCATE (kpp1_env%deriv_set_admm, kpp1_env%rho_set_admm)
               CALL xc_prep_2nd_deriv(kpp1_env%deriv_set_admm, kpp1_env%rho_set_admm, &
                                      rho_r, auxbas_pw_pool, &
                                      xc_section=xc_section_aux)
            END IF
         END IF
      END IF

      !-------------------------------!
      ! Add both hartree and xc terms !
      !-------------------------------!
      DO ispin = 1, nspins
         CALL dbcsr_set(kpp1_env%v_ao(ispin)%matrix, 0.0_dp)

         IF (gapw_xc) THEN
            ! XC and Hartree are integrated separatedly
            ! XC uses the soft basis set only

            IF (nspins == 1) THEN

               IF (.NOT. (lr_triplet)) THEN
                  CALL pw_scale(v_rspace_new(1), 2.0_dp)
                  IF (ASSOCIATED(v_xc_tau)) CALL pw_scale(v_xc_tau(1), 2.0_dp)
               END IF
               CALL qs_rho_get(rho1, rho_ao=rho1_ao)
               ! remove kpp1_env%v_ao and work directly on k_p_p1 ?
               CALL integrate_v_rspace(v_rspace=v_rspace_new(ispin), &
                                       pmat=rho1_ao(ispin), &
                                       hmat=kpp1_env%v_ao(ispin), &
                                       qs_env=qs_env, &
                                       calculate_forces=.FALSE., gapw=gapw_xc)

               IF (ASSOCIATED(v_xc_tau)) THEN
                  CALL integrate_v_rspace(v_rspace=v_xc_tau(ispin), &
                                          pmat=rho1_ao(ispin), &
                                          hmat=kpp1_env%v_ao(ispin), &
                                          qs_env=qs_env, &
                                          compute_tau=.TRUE., &
                                          calculate_forces=.FALSE., gapw=gapw_xc)
               END IF

               ! add hartree only for SINGLETS
               IF (.NOT. lr_triplet) THEN
                  CALL pw_axpy(v_hartree_rspace, v_rspace_new(1), 2.0_dp, 0.0_dp)

                  CALL integrate_v_rspace(v_rspace=v_rspace_new(ispin), &
                                          pmat=rho_ao(ispin), &
                                          hmat=kpp1_env%v_ao(ispin), &
                                          qs_env=qs_env, &
                                          calculate_forces=.FALSE., gapw=gapw)
               END IF
            ELSE
               ! remove kpp1_env%v_ao and work directly on k_p_p1 ?
               CALL integrate_v_rspace(v_rspace=v_rspace_new(ispin), &
                                       pmat=rho_ao(ispin), &
                                       hmat=kpp1_env%v_ao(ispin), &
                                       qs_env=qs_env, &
                                       calculate_forces=.FALSE., gapw=gapw_xc)

               IF (ASSOCIATED(v_xc_tau)) THEN
                  CALL integrate_v_rspace(v_rspace=v_xc_tau(ispin), &
                                          pmat=rho_ao(ispin), &
                                          hmat=kpp1_env%v_ao(ispin), &
                                          qs_env=qs_env, &
                                          compute_tau=.TRUE., &
                                          calculate_forces=.FALSE., gapw=gapw_xc)
               END IF

               CALL pw_copy(v_hartree_rspace, v_rspace_new(ispin))
               CALL integrate_v_rspace(v_rspace=v_rspace_new(ispin), &
                                       pmat=rho_ao(ispin), &
                                       hmat=kpp1_env%v_ao(ispin), &
                                       qs_env=qs_env, &
                                       calculate_forces=.FALSE., gapw=gapw)
            END IF

         ELSE

            IF (nspins == 1) THEN
               IF (.NOT. (lr_triplet)) THEN
                  CALL pw_scale(v_rspace_new(1), 2.0_dp)
                  IF (ASSOCIATED(v_xc_tau)) CALL pw_scale(v_xc_tau(1), 2.0_dp)
               END IF
               ! add hartree only for SINGLETS
               !IF (res_etype == tddfpt_singlet) THEN
               IF (.NOT. lr_triplet) THEN
                  CALL pw_axpy(v_hartree_rspace, v_rspace_new(1), 2.0_dp)
               END IF
            ELSE
               CALL pw_axpy(v_hartree_rspace, v_rspace_new(ispin), 1.0_dp)
            END IF

            IF (lrigpw) THEN
               IF (ASSOCIATED(v_xc_tau)) &
                  CPABORT("metaGGA-functionals not supported with LRI!")

               lri_v_int => lri_density%lri_coefs(ispin)%lri_kinds
               CALL get_qs_env(qs_env, nkind=nkind)
               DO ikind = 1, nkind
                  lri_v_int(ikind)%v_int = 0.0_dp
               END DO
               CALL integrate_v_rspace_one_center(v_rspace_new(ispin), qs_env, &
                                                  lri_v_int, .FALSE., "LRI_AUX")
               DO ikind = 1, nkind
                  CALL para_env%sum(lri_v_int(ikind)%v_int)
               END DO
               ALLOCATE (k1mat(1))
               k1mat(1)%matrix => kpp1_env%v_ao(ispin)%matrix
               IF (lri_env%exact_1c_terms) THEN
                  CALL integrate_v_rspace_diagonal(v_rspace_new(ispin), k1mat(1)%matrix, &
                                                   rho_ao(ispin)%matrix, qs_env, .FALSE., "ORB")
               END IF
               CALL calculate_lri_ks_matrix(lri_env, lri_v_int, k1mat, atomic_kind_set)
               DEALLOCATE (k1mat)
            ELSE
               CALL integrate_v_rspace(v_rspace=v_rspace_new(ispin), &
                                       pmat=rho_ao(ispin), &
                                       hmat=kpp1_env%v_ao(ispin), &
                                       qs_env=qs_env, &
                                       calculate_forces=.FALSE., gapw=gapw)

               IF (ASSOCIATED(v_xc_tau)) THEN
                  CALL integrate_v_rspace(v_rspace=v_xc_tau(ispin), &
                                          pmat=rho_ao(ispin), &
                                          hmat=kpp1_env%v_ao(ispin), &
                                          qs_env=qs_env, &
                                          compute_tau=.TRUE., &
                                          calculate_forces=.FALSE., gapw=gapw)
               END IF
            END IF

         END IF

         CALL dbcsr_copy(p_env%kpp1(ispin)%matrix, kpp1_env%v_ao(ispin)%matrix)
      END DO

      IF (gapw) THEN
         IF (.NOT. ((nspins == 1 .AND. lr_triplet))) THEN
            CALL Vh_1c_gg_integrals(qs_env, energy_hartree_1c, &
                                    p_env%hartree_local%ecoul_1c, &
                                    p_env%local_rho_set, &
                                    para_env, tddft=.TRUE., core_2nd=.TRUE.)

            CALL integrate_vhg0_rspace(qs_env, v_hartree_rspace, para_env, &
                                       calculate_forces=.FALSE., &
                                       local_rho_set=p_env%local_rho_set)
         END IF
         ! ***  Add single atom contributions to the KS matrix ***
         ! remap pointer
         ns = SIZE(p_env%kpp1)
         ksmat(1:ns, 1:1) => p_env%kpp1(1:ns)
         ns = SIZE(rho_ao)
         psmat(1:ns, 1:1) => rho_ao(1:ns)
         CALL update_ks_atom(qs_env, ksmat, psmat, forces=.FALSE., tddft=.TRUE., &
                             rho_atom_external=p_env%local_rho_set%rho_atom_set)
      ELSEIF (gapw_xc) THEN
         ns = SIZE(p_env%kpp1)
         ksmat(1:ns, 1:1) => p_env%kpp1(1:ns)
         ns = SIZE(rho_ao)
         psmat(1:ns, 1:1) => rho_ao(1:ns)
         CALL update_ks_atom(qs_env, ksmat, psmat, forces=.FALSE., tddft=.TRUE., &
                             rho_atom_external=p_env%local_rho_set%rho_atom_set)
      END IF

      ! KG embedding, contribution of kinetic energy functional to kernel
      IF (dft_control%qs_control%do_kg .AND. .NOT. (lr_triplet .OR. gapw .OR. gapw_xc)) THEN
         IF (qs_env%kg_env%tnadd_method == kg_tnadd_embed) THEN

            CALL qs_rho_get(rho1, rho_ao=rho1_ao)
            alpha = 1.0_dp

            ekin_mol = 0.0_dp
            CALL get_qs_env(qs_env, kg_env=kg_env)
            CALL kg_ekin_subset(qs_env=qs_env, &
                                ks_matrix=p_env%kpp1, &
                                ekin_mol=ekin_mol, &
                                calc_force=.FALSE., &
                                do_kernel=.TRUE., &
                                pmat_ext=rho1_ao)
         END IF
      END IF

      CALL auxbas_pw_pool%give_back_pw(v_hartree_gspace)
      CALL auxbas_pw_pool%give_back_pw(v_hartree_rspace)
      DO ispin = 1, nspins
         CALL auxbas_pw_pool%give_back_pw(v_rspace_new(ispin))
      END DO
      DEALLOCATE (v_rspace_new)
      IF (ASSOCIATED(v_xc_tau)) THEN
         DO ispin = 1, nspins
            CALL auxbas_pw_pool%give_back_pw(v_xc_tau(ispin))
         END DO
         DEALLOCATE (v_xc_tau)
      END IF

      CALL timestop(handle)

   END SUBROUTINE apply_op_2_dft

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param p_env ...
! **************************************************************************************************
   SUBROUTINE apply_op_2_xtb(qs_env, p_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_p_env_type)                                :: p_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'apply_op_2_xtb'

      INTEGER                                            :: atom_a, handle, iatom, ikind, is, ispin, &
                                                            na, natom, natorb, nkind, ns, nsgf, &
                                                            nspins
      INTEGER, DIMENSION(25)                             :: lao
      INTEGER, DIMENSION(5)                              :: occ
      LOGICAL                                            :: lr_triplet
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: mcharge, mcharge1
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: aocg, aocg1, charges, charges1
      TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: rho_ao
      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: matrix_p, matrix_p1, matrix_s
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(linres_control_type), POINTER                 :: linres_control
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(qs_kpp1_env_type), POINTER                    :: kpp1_env
      TYPE(qs_rho_type), POINTER                         :: rho, rho1
      TYPE(xtb_atom_type), POINTER                       :: xtb_kind

      CALL timeset(routineN, handle)

      CPASSERT(ASSOCIATED(p_env%kpp1_env))
      CPASSERT(ASSOCIATED(p_env%kpp1))
      kpp1_env => p_env%kpp1_env

      rho1 => p_env%rho1
      CPASSERT(ASSOCIATED(rho1))

      CALL get_qs_env(qs_env=qs_env, &
                      pw_env=pw_env, &
                      para_env=para_env, &
                      rho=rho, &
                      linres_control=linres_control, &
                      dft_control=dft_control)

      CALL qs_rho_get(rho, rho_ao=rho_ao)

      lr_triplet = linres_control%lr_triplet
      CPASSERT(.NOT. lr_triplet)

      nspins = SIZE(p_env%kpp1)

      DO ispin = 1, nspins
         CALL dbcsr_set(p_env%kpp1(ispin)%matrix, 0.0_dp)
      END DO

      IF (dft_control%qs_control%xtb_control%coulomb_interaction) THEN
         ! Mulliken charges
         CALL get_qs_env(qs_env, particle_set=particle_set, matrix_s_kp=matrix_s)
         natom = SIZE(particle_set)
         CALL qs_rho_get(rho, rho_ao_kp=matrix_p)
         CALL qs_rho_get(rho1, rho_ao_kp=matrix_p1)
         ALLOCATE (mcharge(natom), charges(natom, 5))
         ALLOCATE (mcharge1(natom), charges1(natom, 5))
         charges = 0.0_dp
         charges1 = 0.0_dp
         CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set, qs_kind_set=qs_kind_set)
         nkind = SIZE(atomic_kind_set)
         CALL get_qs_kind_set(qs_kind_set, maxsgf=nsgf)
         ALLOCATE (aocg(nsgf, natom))
         aocg = 0.0_dp
         ALLOCATE (aocg1(nsgf, natom))
         aocg1 = 0.0_dp
         CALL ao_charges(matrix_p, matrix_s, aocg, para_env)
         CALL ao_charges(matrix_p1, matrix_s, aocg1, para_env)
         DO ikind = 1, nkind
            CALL get_atomic_kind(atomic_kind_set(ikind), natom=na)
            CALL get_qs_kind(qs_kind_set(ikind), xtb_parameter=xtb_kind)
            CALL get_xtb_atom_param(xtb_kind, natorb=natorb, lao=lao, occupation=occ)
            DO iatom = 1, na
               atom_a = atomic_kind_set(ikind)%atom_list(iatom)
               charges(atom_a, :) = REAL(occ(:), KIND=dp)
               DO is = 1, natorb
                  ns = lao(is) + 1
                  charges(atom_a, ns) = charges(atom_a, ns) - aocg(is, atom_a)
                  charges1(atom_a, ns) = charges1(atom_a, ns) - aocg1(is, atom_a)
               END DO
            END DO
         END DO
         DEALLOCATE (aocg, aocg1)
         DO iatom = 1, natom
            mcharge(iatom) = SUM(charges(iatom, :))
            mcharge1(iatom) = SUM(charges1(iatom, :))
         END DO
         ! Coulomb Kernel
         CALL xtb_coulomb_hessian(qs_env, p_env%kpp1, charges1, mcharge1, mcharge)
         !
         DEALLOCATE (charges, mcharge, charges1, mcharge1)
      END IF

      CALL timestop(handle)

   END SUBROUTINE apply_op_2_xtb

! **************************************************************************************************
!> \brief Update action of TDDFPT operator on trial vectors by adding exact-exchange term.
!> \param qs_env ...
!> \param p_env ...
!> \par History
!>    * 11.2019 adapted from tddfpt_apply_hfx
! **************************************************************************************************
   SUBROUTINE apply_hfx(qs_env, p_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_p_env_type)                                :: p_env

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'apply_hfx'

      INTEGER                                            :: handle, ispin, nspins
      LOGICAL                                            :: do_hfx
      REAL(KIND=dp)                                      :: alpha
      TYPE(cp_logger_type), POINTER                      :: logger
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: h1_mat, matrix_s, rho1_ao, work
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(section_vals_type), POINTER                   :: hfx_section, input

      CALL timeset(routineN, handle)

      logger => cp_get_default_logger()

      CALL get_qs_env(qs_env=qs_env, &
                      input=input, &
                      matrix_s=matrix_s, &
                      dft_control=dft_control)
      nspins = dft_control%nspins

      hfx_section => section_vals_get_subs_vals(input, "DFT%XC%HF")
      CALL section_vals_get(hfx_section, explicit=do_hfx)

      IF (do_hfx) THEN

         IF (dft_control%do_admm) THEN
            IF (dft_control%admm_control%purification_method /= do_admm_purify_none) THEN
               CPABORT("ADMM: Linear Response needs purification_method=none")
            END IF
            IF (dft_control%admm_control%scaling_model /= do_admm_exch_scaling_none) THEN
               CPABORT("ADMM: Linear Response needs scaling_model=none")
            END IF
            IF (dft_control%admm_control%method /= do_admm_basis_projection) THEN
               CPABORT("ADMM: Linear Response needs admm_method=basis_projection")
            END IF
            !
            rho1_ao => p_env%p1_admm
            h1_mat => p_env%kpp1_admm
         ELSE
            rho1_ao => p_env%p1
            h1_mat => p_env%kpp1
         END IF

         NULLIFY (work)
         CALL dbcsr_allocate_matrix_set(work, nspins)
         DO ispin = 1, nspins
            ALLOCATE (work(ispin)%matrix)
            CALL dbcsr_create(work(ispin)%matrix, template=h1_mat(ispin)%matrix)
            CALL dbcsr_copy(work(ispin)%matrix, h1_mat(ispin)%matrix)
            CALL dbcsr_set(work(ispin)%matrix, 0.0_dp)
         END DO

         CALL hfx_matrix(work, rho1_ao, qs_env, hfx_section)

         alpha = 2.0_dp
         IF (nspins == 2) alpha = 1.0_dp

         DO ispin = 1, nspins
            CALL dbcsr_add(h1_mat(ispin)%matrix, work(ispin)%matrix, 1.0_dp, alpha)
         END DO

         CALL dbcsr_deallocate_matrix_set(work)

      END IF

      CALL timestop(handle)

   END SUBROUTINE apply_hfx

! **************************************************************************************************
!> \brief Add the hfx contributions to the Hamiltonian
!>
!> \param matrix_ks ...
!> \param rho_ao ...
!> \param qs_env ...
!> \param hfx_sections ...
!> \param external_x_data ...
!> \param ex ...
!> \note
!>     Simplified version of subroutine hfx_ks_matrix()
! **************************************************************************************************
   SUBROUTINE hfx_matrix(matrix_ks, rho_ao, qs_env, hfx_sections, external_x_data, ex)
      TYPE(dbcsr_p_type), DIMENSION(:), TARGET           :: matrix_ks, rho_ao
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(section_vals_type), POINTER                   :: hfx_sections
      TYPE(hfx_type), DIMENSION(:, :), OPTIONAL, TARGET  :: external_x_data
      REAL(KIND=dp), OPTIONAL                            :: ex

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'hfx_matrix'

      INTEGER                                            :: handle, irep, ispin, mspin, n_rep_hf, &
                                                            nspins
      LOGICAL                                            :: distribute_fock_matrix, &
                                                            hfx_treat_lsd_in_core, &
                                                            s_mstruct_changed
      REAL(KIND=dp)                                      :: eh1, ehfx
      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: matrix_ks_kp, rho_ao_kp
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(hfx_type), DIMENSION(:, :), POINTER           :: x_data
      TYPE(mp_para_env_type), POINTER                    :: para_env

      CALL timeset(routineN, handle)

      NULLIFY (dft_control, para_env, matrix_ks_kp, rho_ao_kp, x_data)

      CALL get_qs_env(qs_env=qs_env, &
                      dft_control=dft_control, &
                      para_env=para_env, &
                      s_mstruct_changed=s_mstruct_changed, &
                      x_data=x_data)

      IF (PRESENT(external_x_data)) x_data => external_x_data

      CPASSERT(dft_control%nimages == 1)
      nspins = dft_control%nspins

      CALL section_vals_get(hfx_sections, n_repetition=n_rep_hf)
      CALL section_vals_val_get(hfx_sections, "TREAT_LSD_IN_CORE", l_val=hfx_treat_lsd_in_core, &
                                i_rep_section=1)

      CALL section_vals_get(hfx_sections, n_repetition=n_rep_hf)
      distribute_fock_matrix = .TRUE.

      mspin = 1
      IF (hfx_treat_lsd_in_core) mspin = nspins

      matrix_ks_kp(1:nspins, 1:1) => matrix_ks(1:nspins)
      rho_ao_kp(1:nspins, 1:1) => rho_ao(1:nspins)

      DO irep = 1, n_rep_hf
         ehfx = 0.0_dp

         IF (x_data(irep, 1)%do_hfx_ri) THEN
            CALL hfx_ri_update_ks(qs_env, x_data(irep, 1)%ri_data, matrix_ks_kp, ehfx, &
                                  rho_ao=rho_ao_kp, geometry_did_change=s_mstruct_changed, &
                                  nspins=nspins, hf_fraction=x_data(irep, 1)%general_parameter%fraction)

         ELSE

            DO ispin = 1, mspin
               CALL integrate_four_center(qs_env, x_data, matrix_ks_kp, eh1, rho_ao_kp, hfx_sections, para_env, &
                                          s_mstruct_changed, irep, distribute_fock_matrix, ispin=ispin)
               ehfx = ehfx + eh1
            END DO

         END IF
      END DO

      ! Export energy
      IF (PRESENT(ex)) ex = ehfx

      CALL timestop(handle)

   END SUBROUTINE hfx_matrix

! **************************************************************************************************
!> \brief ...
!> \param qs_env ...
!> \param p_env ...
! **************************************************************************************************
   SUBROUTINE apply_xc_admm(qs_env, p_env)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_p_env_type)                                :: p_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'apply_xc_admm'

      CHARACTER(LEN=default_string_length)               :: basis_type
      INTEGER                                            :: handle, ispin, ns, nspins
      INTEGER, DIMENSION(2, 3)                           :: bo
      LOGICAL                                            :: gapw, lsd
      REAL(KIND=dp)                                      :: alpha
      TYPE(admm_type), POINTER                           :: admm_env
      TYPE(dbcsr_p_type)                                 :: xcmat
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_s
      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: ksmat, psmat
      TYPE(dft_control_type), POINTER                    :: dft_control
      TYPE(linres_control_type), POINTER                 :: linres_control
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(neighbor_list_set_p_type), DIMENSION(:), &
         POINTER                                         :: sab_aux_fit
      TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER        :: rho1_aux_g
      TYPE(pw_env_type), POINTER                         :: pw_env
      TYPE(pw_pool_type), POINTER                        :: auxbas_pw_pool
      TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER        :: rho1_aux_r, tau_pw, v_xc, v_xc_tau
      TYPE(rho_atom_type), DIMENSION(:), POINTER         :: rho1_atom_set, rho_atom_set
      TYPE(section_vals_type), POINTER                   :: xc_fun_section, xc_section
      TYPE(task_list_type), POINTER                      :: task_list
      TYPE(xc_rho_cflags_type)                           :: needs
      TYPE(xc_rho_set_type)                              :: rho1_set

      CALL timeset(routineN, handle)

      CALL get_qs_env(qs_env=qs_env, dft_control=dft_control)

      IF (dft_control%do_admm) THEN
         IF (qs_env%admm_env%aux_exch_func == do_admm_aux_exch_func_none) THEN
            ! nothing to do
         ELSE
            CALL get_qs_env(qs_env=qs_env, linres_control=linres_control)
            CPASSERT(.NOT. dft_control%qs_control%lrigpw)
            CPASSERT(.NOT. linres_control%lr_triplet)

            nspins = dft_control%nspins

            ! AUX basis contribution
            CALL get_qs_env(qs_env=qs_env, pw_env=pw_env)
            CPASSERT(ASSOCIATED(pw_env))
            CALL pw_env_get(pw_env, auxbas_pw_pool=auxbas_pw_pool)
            NULLIFY (tau_pw)
            ! calculate the xc potential
            lsd = (nspins == 2)
            CALL get_admm_env(qs_env%admm_env, matrix_s_aux_fit=matrix_s)
            ALLOCATE (xcmat%matrix)
            CALL dbcsr_create(xcmat%matrix, template=matrix_s(1)%matrix)

            CALL get_qs_env(qs_env, admm_env=admm_env)
            gapw = admm_env%do_gapw

            CALL qs_rho_get(p_env%rho1_admm, rho_r=rho1_aux_r, rho_g=rho1_aux_g)
            xc_section => admm_env%xc_section_aux
            bo = rho1_aux_r(1)%pw_grid%bounds_local
            ! create the place where to store the argument for the functionals
            CALL xc_rho_set_create(rho1_set, bo, &
                                   rho_cutoff=section_get_rval(xc_section, "DENSITY_CUTOFF"), &
                                   drho_cutoff=section_get_rval(xc_section, "GRADIENT_CUTOFF"), &
                                   tau_cutoff=section_get_rval(xc_section, "TAU_CUTOFF"))

            xc_fun_section => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
            needs = xc_functionals_get_needs(xc_fun_section, lsd, .TRUE.)

            ! calculate the arguments needed by the functionals
            CALL xc_rho_set_update(rho1_set, rho1_aux_r, rho1_aux_g, tau_pw, needs, &
                                   section_get_ival(xc_section, "XC_GRID%XC_DERIV"), &
                                   section_get_ival(xc_section, "XC_GRID%XC_SMOOTH_RHO"), &
                                   auxbas_pw_pool)
            CALL xc_calc_2nd_deriv(v_xc, v_xc_tau, p_env%kpp1_env%deriv_set_admm, p_env%kpp1_env%rho_set_admm, &
                                   rho1_aux_r, rho1_aux_g, tau_pw, auxbas_pw_pool, gapw=.FALSE., &
                                   xc_section=xc_section)
            IF (ASSOCIATED(v_xc_tau)) THEN
               CPABORT("Meta-GGA ADMM functionals not yet supported!")
            END IF
            CALL xc_rho_set_release(rho1_set)

            basis_type = "AUX_FIT"
            CALL get_qs_env(qs_env, para_env=para_env)
            CALL get_admm_env(admm_env, task_list_aux_fit=task_list)
            IF (admm_env%do_gapw) THEN
               CALL prepare_gapw_den(qs_env, local_rho_set=p_env%local_rho_set_admm, &
                                     do_rho0=.FALSE., kind_set_external=admm_env%admm_gapw_env%admm_kind_set)
               rho_atom_set => admm_env%admm_gapw_env%local_rho_set%rho_atom_set
               rho1_atom_set => p_env%local_rho_set_admm%rho_atom_set
               CALL calculate_xc_2nd_deriv_atom(rho_atom_set, rho1_atom_set, qs_env, xc_section, para_env, &
                                                kind_set_external=admm_env%admm_gapw_env%admm_kind_set)
               basis_type = "AUX_FIT_SOFT"
               task_list => admm_env%admm_gapw_env%task_list
            END IF

            alpha = 1.0_dp
            IF (nspins == 1) alpha = 2.0_dp

            DO ispin = 1, nspins
               CALL pw_scale(v_xc(ispin), v_xc(ispin)%pw_grid%dvol)
               CALL dbcsr_copy(xcmat%matrix, matrix_s(1)%matrix)
               CALL dbcsr_set(xcmat%matrix, 0.0_dp)
               CALL integrate_v_rspace(v_rspace=v_xc(ispin), hmat=xcmat, qs_env=qs_env, &
                                       calculate_forces=.FALSE., basis_type=basis_type, &
                                       task_list_external=task_list)
               CALL dbcsr_add(p_env%kpp1_admm(ispin)%matrix, xcmat%matrix, 1.0_dp, alpha)
            END DO

            IF (admm_env%do_gapw) THEN
               CALL get_admm_env(admm_env, sab_aux_fit=sab_aux_fit)
               ns = SIZE(p_env%kpp1_admm)
               ksmat(1:ns, 1:1) => p_env%kpp1_admm(1:ns)
               psmat(1:ns, 1:1) => p_env%p1_admm(1:ns)
               CALL update_ks_atom(qs_env, ksmat, psmat, forces=.FALSE., tddft=.TRUE., &
                                   rho_atom_external=p_env%local_rho_set_admm%rho_atom_set, &
                                   kind_set_external=admm_env%admm_gapw_env%admm_kind_set, &
                                   oce_external=admm_env%admm_gapw_env%oce, &
                                   sab_external=sab_aux_fit)
            END IF

            DO ispin = 1, nspins
               CALL auxbas_pw_pool%give_back_pw(v_xc(ispin))
            END DO
            DEALLOCATE (v_xc)
            CALL dbcsr_deallocate_matrix(xcmat%matrix)

         END IF
      END IF

      CALL timestop(handle)

   END SUBROUTINE apply_xc_admm

END MODULE qs_linres_kernel
